#include <iostream>
#include <cstdio>
#include <helper_thread_idx.h>

__global__ void test_kernel(int *out) {
    int tid = getGlobalIdx_1D_1D();
    out[tid] = tid;
}

// Test function for 1D grid of 1D blocks
void test_1D_1D() {
    const int N = 1024;
    const int BLOCK_SIZE = 128;
    const int GRID_SIZE = (N + BLOCK_SIZE - 1) / BLOCK_SIZE;

    int *d_out, *h_out;
    cudaMalloc(&d_out, N * sizeof(int));
    h_out = new int[N];

    // Launch kernel with 1D grid of 1D blocks
    test_kernel<<<GRID_SIZE, BLOCK_SIZE>>>(d_out);

    cudaMemcpy(h_out, d_out, N * sizeof(int), cudaMemcpyDeviceToHost);

    // Verify output
    for (int i = 0; i < N; i++) {
        if (h_out[i] != i) {
            printf("ERROR: expected value %d at index %d, but got %d\n", i, i, h_out[i]);
            break;
        }
    }

    cudaFree(d_out);
    delete[] h_out;
}

// Test function for 1D grid of 2D blocks
// void test_1D_2D() {
//     const int N = 1024;
//     const int BLOCK_SIZE_X = 16;
//     const int BLOCK_SIZE_Y = 8;
//     const int GRID_SIZE = (N + (BLOCK_SIZE_X * BLOCK_SIZE_Y) - 1) / (BLOCK_SIZE_X * BLOCK_SIZE_Y);

//     int *d_out, *h_out;
//     cudaMalloc(&d_out, N * sizeof(int));
//     h_out = new int[N];

//     dim3 block_dim(BLOCK_SIZE_X, BLOCK_SIZE_Y);
//     // Launch kernel with 1D grid of 2D blocks
//     test_kernel<<<GRID_SIZE, block_dim>>>(d_out, getGlobalIdx_1D_2D);

//     cudaMemcpy(h_out, d_out, N * sizeof(int), cudaMemcpyDeviceToHost);

//     // Verify output
//     for (int i = 0; i < N; i++) {
//         if (h_out[i] != i) {
//             printf("ERROR: expected value %d at index %d, but got %d\n", i, i, h_out[i]);
//             break;
//         }
//     }

//     cudaFree(d_out);
//     delete[] h_out;
// }

// // Test function for 1D grid of 3D blocks
// void test_1D_3D() {
//     const int N = 1024;
//     const int BLOCK_SIZE_X = 4;
//     const int BLOCK_SIZE_Y = 4;
//     const int BLOCK_SIZE_Z = 4;
//     const int GRID_SIZE = (N + (BLOCK_SIZE_X * BLOCK_SIZE_Y * BLOCK_SIZE_Z) - 1) / (BLOCK_SIZE_X * BLOCK_SIZE_Y * BLOCK_SIZE_Z);

//     int *d_out, *h_out;
//     cudaMalloc(&d_out, N * sizeof(int));
//     h_out = new int[N];

//     dim3 block_dim(BLOCK_SIZE_X, BLOCK_SIZE_Y, BLOCK_SIZE_Z);
//     // Launch kernel with 1D grid of 3D blocks
//     test_kernel<<<GRID_SIZE, block_dim>>>(d_out, getGlobalIdx_1D_3D);

//     cudaMemcpy(h_out, d_out, N * sizeof(int), cudaMemcpyDeviceToHost);

//     // Verify output
//     for (int i = 0; i < N; i++) {
//         if (h_out[i] != i) {
//             printf("ERROR: expected value %d at index %d, but got %d\n", i, i, h_out[i]);
//             break;
//         }
//     }

//     cudaFree(d_out);
//     delete[] h_out;
// }

// // Test function for 2D grid of 1D blocks
// void test_2D_1D() {
//     const int N = 1024;
//     const int BLOCK_SIZE = 64;
//     const int GRID_SIZE_X = 16;
//     const int GRID_SIZE_Y = (N + (BLOCK_SIZE * GRID_SIZE_X) - 1) / (BLOCK_SIZE * GRID_SIZE_X);

//     int *d_out, *h_out;
//     cudaMalloc(&d_out, N * sizeof(int));
//     h_out = new int[N];

//     dim3 grid_dim(GRID_SIZE_X, GRID_SIZE_Y);
//     // Launch kernel with 2D grid of 1D blocks
//     test_kernel<<<grid_dim, BLOCK_SIZE>>>(d_out, getGlobalIdx_2D_1D);

//     cudaMemcpy(h_out, d_out, N * sizeof(int), cudaMemcpyDeviceToHost);

//     // Verify output
//     for (int i = 0; i < N; i++) {
//         if (h_out[i] != i) {
//             printf("ERROR: expected value %d at index %d, but got %d\n", i, i, h_out[i]);
//             break;
//         }
//     }

//     cudaFree(d_out);
//     delete[] h_out;
// }

// // Test function for 2D grid of 2D blocks
// void test_2D_2D() {
//     const int N = 1024;
//     const int BLOCK_SIZE_X = 8;
//     const int BLOCK_SIZE_Y = 8;
//     const int GRID_SIZE_X = (N + (BLOCK_SIZE_X * BLOCK_SIZE_Y) - 1) / (BLOCK_SIZE_X * BLOCK_SIZE_Y);
//     const int GRID_SIZE_Y = (N + (BLOCK_SIZE_X * BLOCK_SIZE_Y * GRID_SIZE_X) - 1) / (BLOCK_SIZE_X * BLOCK_SIZE_Y * GRID_SIZE_X);

//     int *d_out, *h_out;
//     cudaMalloc(&d_out, N * sizeof(int));
//     h_out = new int[N];

//     dim3 grid_dim(GRID_SIZE_X, GRID_SIZE_Y);
//     dim3 block_dim(BLOCK_SIZE_X, BLOCK_SIZE_Y);
//     // Launch kernel with 2D grid of 2D blocks
//     test_kernel<<<grid_dim, block_dim>>>(d_out, getGlobalIdx_2D_2D);

//     cudaMemcpy(h_out, d_out, N * sizeof(int), cudaMemcpyDeviceToHost);

//     // Verify output
//     for (int i = 0; i < N; i++) {
//         if (h_out[i] != i) {
//             printf("ERROR: expected value %d at index %d, but got %d\n", i, i, h_out[i]);
//             break;
//         }
//     }

//     cudaFree(d_out);
//     delete[] h_out;
// }

// // Test function for 2D grid of 3D blocks
// void test_2D_3D() {
//     const int N = 1024;
//     const int BLOCK_SIZE_X = 4;
//     const int BLOCK_SIZE_Y = 4;
//     const int BLOCK_SIZE_Z = 4;
//     const int GRID_SIZE_X = (N + (BLOCK_SIZE_X * BLOCK_SIZE_Y * BLOCK_SIZE_Z) - 1) / (BLOCK_SIZE_X * BLOCK_SIZE_Y * BLOCK_SIZE_Z);
//     const int GRID_SIZE_Y = (N + (BLOCK_SIZE_X * BLOCK_SIZE_Y * BLOCK_SIZE_Z * GRID_SIZE_X) - 1) / (BLOCK_SIZE_X * BLOCK_SIZE_Y * BLOCK_SIZE_Z * GRID_SIZE_X);

//     int *d_out, *h_out;
//     cudaMalloc(&d_out, N * sizeof(int));
//     h_out = new int[N];

//     dim3 grid_dim(GRID_SIZE_X, GRID_SIZE_Y);
//     dim3 block_dim(BLOCK_SIZE_X, BLOCK_SIZE_Y, BLOCK_SIZE_Z);
//     // Launch kernel with 2D grid of 3D blocks
//     test_kernel<<<grid_dim, block_dim>>>(d_out, getGlobalIdx_2D_3D);

//     cudaMemcpy(h_out, d_out, N * sizeof(int), cudaMemcpyDeviceToHost);

//     // Verify output
//     for (int i = 0; i < N; i++) {
//         if (h_out[i] != i) {
//             printf("ERROR: expected value %d at index %d, but got %d\n", i, i, h_out[i]);
//             break;
//         }
//     }

//     cudaFree(d_out);
//     delete[] h_out;
// }

// // Test function for 3D grid of 1D blocks
// void test_3D_1D() {
//     const int N = 1024;
//     const int BLOCK_SIZE = 8;
//     const int GRID_SIZE_X = 4;
//     const int GRID_SIZE_Y = 4;
//     const int GRID_SIZE_Z = (N + (BLOCK_SIZE * GRID_SIZE_X * GRID_SIZE_Y) - 1) / (BLOCK_SIZE * GRID_SIZE_X * GRID_SIZE_Y);

//     int *d_out, *h_out;
//     cudaMalloc(&d_out, N * sizeof(int));
//     h_out = new int[N];

//     dim3 grid_dim(GRID_SIZE_X, GRID_SIZE_Y, GRID_SIZE_Z);
//     // Launch kernel with 3D grid of 1D blocks
//     test_kernel<<<grid_dim, BLOCK_SIZE>>>(d_out, getGlobalIdx_3D_1D);

//     cudaMemcpy(h_out, d_out, N * sizeof(int), cudaMemcpyDeviceToHost);

//     // Verify output
//     for (int i = 0; i < N; i++) {
//         if (h_out[i] != i) {
//             printf("ERROR: expected value %d at index %d, but got %d\n", i, i, h_out[i]);
//             break;
//         }
//     }

//     cudaFree(d_out);
//     delete[] h_out;
// }

// // Test function for 3D grid of 2D blocks
// void test_3D_2D() {
//     const int N = 1024;
//     const int BLOCK_SIZE_X = 4;
//     const int BLOCK_SIZE_Y = 4;
//     const int BLOCK_SIZE_Z = 4;
//     const int GRID_SIZE_X = (N + (BLOCK_SIZE_X * BLOCK_SIZE_Y * BLOCK_SIZE_Z) - 1) / (BLOCK_SIZE_X * BLOCK_SIZE_Y * BLOCK_SIZE_Z);
//     const int GRID_SIZE_Y = (N + (BLOCK_SIZE_X * BLOCK_SIZE_Y * BLOCK_SIZE_Z * GRID_SIZE_X) - 1) / (BLOCK_SIZE_X * BLOCK_SIZE_Y * BLOCK_SIZE_Z * GRID_SIZE_X);
//     const int GRID_SIZE_Z = (N + (BLOCK_SIZE_X * BLOCK_SIZE_Y * GRID_SIZE_X * GRID_SIZE_Y) - 1) / (BLOCK_SIZE_X * BLOCK_SIZE_Y * GRID_SIZE_X * GRID_SIZE_Y);

//     int *d_out, *h_out;
//     cudaMalloc(&d_out, N * sizeof(int));
//     h_out = new int[N];

//     dim3 grid_dim(GRID_SIZE_X, GRID_SIZE_Y, GRID_SIZE_Z);
//     dim3 block_dim(BLOCK_SIZE_X, BLOCK_SIZE_Y, BLOCK_SIZE_Z);
//     // Launch kernel with 3D grid of 2D blocks
//     test_kernel<<<grid_dim, block_dim>>>(d_out, getGlobalIdx_3D_2D);

//     cudaMemcpy(h_out, d_out, N * sizeof(int), cudaMemcpyDeviceToHost);

//     // Verify output
//     for (int i = 0; i < N; i++) {
//         if (h_out[i] != i) {
//             printf("ERROR: expected value %d at index %d, but got %d\n", i, i, h_out[i]);
//             break;
//         }
//     }

//     cudaFree(d_out);
//     delete[] h_out;
// }

// // Test function for 3D grid of 3D blocks
// void test_3D_3D() {
//     const int N = 1024;
//     const int BLOCK_SIZE_X = 2;
//     const int BLOCK_SIZE_Y = 2;
//     const int BLOCK_SIZE_Z = 2;
//     const int GRID_SIZE_X = (N + (BLOCK_SIZE_X * BLOCK_SIZE_Y * BLOCK_SIZE_Z) - 1) / (BLOCK_SIZE_X * BLOCK_SIZE_Y * BLOCK_SIZE_Z);
//     const int GRID_SIZE_Y = (N + (BLOCK_SIZE_X * BLOCK_SIZE_Y * BLOCK_SIZE_Z * GRID_SIZE_X) - 1) / (BLOCK_SIZE_X * BLOCK_SIZE_Y * BLOCK_SIZE_Z * GRID_SIZE_X);
//     const int GRID_SIZE_Z = (N + (BLOCK_SIZE_X * BLOCK_SIZE_Y * BLOCK_SIZE_Z * GRID_SIZE_X * GRID_SIZE_Y) - 1) / (BLOCK_SIZE_X * BLOCK_SIZE_Y * BLOCK_SIZE_Z * GRID_SIZE_X * GRID_SIZE_Y);

//     int *d_out, *h_out;
//     cudaMalloc(&d_out, N * sizeof(int));
//     h_out = new int[N];

//     dim3 grid_dim(GRID_SIZE_X, GRID_SIZE_Y, GRID_SIZE_Z);
//     dim3 block_dim(BLOCK_SIZE_X, BLOCK_SIZE_Y, BLOCK_SIZE_Z);
//     // Launch kernel with 3D grid of 3D blocks
//     test_kernel<<<grid_dim, block_dim>>>(d_out, getGlobalIdx_3D_3D);

//     cudaMemcpy(h_out, d_out, N * sizeof(int), cudaMemcpyDeviceToHost);

//     // Verify output
//     for (int i = 0; i < N; i++) {
//         if (h_out[i] != i) {
//             printf("ERROR: expected value %d at index %d, but got %d\n", i, i, h_out[i]);
//             break;
//         }
//     }

//     cudaFree(d_out);
//     delete[] h_out;
// }

int main(){
    test_1D_1D();   
}